from pydantic import BaseModel, Field
from langchain.agents import create_agent
from langchain.agents.structured_output import ProviderStrategy, ToolStrategy
from langchain.agents.middleware import wrap_model_call, ModelRequest, ModelResponse, dynamic_prompt

from model.Ark import ArkModel
from model.Tongyi import TongyiModel

## 基础示例
def base_agent_ep():
    model = ArkModel(
        temperature=0.2,
    ).model

    agent = create_agent(
        model=model,
        system_prompt='你是一个聊天小助手，我的名字叫小新，请你回答用户的提问。'
    )
    response = agent.invoke({
        "messages": [{"role": "user", "content": "我的名字叫什么?"}]
    })

    print(response)

## 动态替换, 动态替换系统提示词、模型名称等
def dynamic_model():
    @wrap_model_call
    def model_middleware(request: ModelRequest, handler) -> ModelResponse:
        ## 动态替换模型，同时，还可以替换提示词等
        request.model = TongyiModel().model
        request.system_prompt = '你是一个聊天小助手，我的名字叫小新，请你回答用户的提问。'
        return handler(request)

    agent = create_agent(
        model=ArkModel().model,
        middleware=[model_middleware]
    )
    response = agent.batch([
        {
            "messages": [{"role": "user", "content": "你的版本是多少？"}]
        },
        {
            "messages": [{"role": "user", "content": "我叫什么名字？"}]
        }
    ])
    print(response)

## 结合context的动态替换，新版的 context字段 替代了旧版的 configurate。
## 其中, context 中的配置类似于 之前的langgraph 中的 state。context 的主要作用就是可以预定义一些配置，并且可以在工具或其它链路中传递消息
def dynamic_context_model():
    class CustomerContextSchema(BaseModel):
        customer_var: int

    @dynamic_prompt
    def middleware_func(request: ModelRequest) -> str:
        print(request.runtime.context['customer_var'])

        if request.runtime.context['customer_var'] == 1:
            return "我的名字叫小新，你叫什么名字？"

        return '我的名字叫小白，你叫什么名字？?'

    agent = create_agent(
        model=ArkModel().model,
        middleware=[middleware_func],
        context_schema=CustomerContextSchema
    )
    response = agent.invoke({
        "messages": [
            { "role": 'user', "content": "我叫什么名字?" }
        ]
    }, context={"customer_var": -1})
    print(response)

## 结构化输出，可定义输出的内容
def output_format():
    class OutputSchema(BaseModel):
        type: str = Field(description='笑话的类型')
        content: str = Field(description='笑话的内容')

    agent = create_agent(
        model=ArkModel().model,
        response_format=ToolStrategy(OutputSchema)
    )

    response = agent.invoke({ 'messages': [{ "role": 'user', "content": "请给我讲一个冷笑话" }] })
    print(response['messages'][-1])

def run():
    output_format()